import pymysql
from datetime import datetime

class DbOprate():
    # MySQL数据库配置
    DB_CONFIG = {
        'host': '127.0.0.1',
        'user': 'root',
        'password': 'root',
        'database': 'stock_db',
        'charset': 'utf8mb4'
    }
    
    '''
    DB_CONFIG = {
        'host': '172.30.168.43',         # 数据库主机地址
        'user': 'root',     # 数据库用户名
        'password': 'DAzA8nuLfDCQfD2n', # 数据库密码
        'database': 'stock_db', # 数据库名称
        'charset': 'utf8mb4'
    }
    '''
    @staticmethod
    def get_connection():
        """获取数据库连接"""
        return pymysql.connect(**DbOprate.DB_CONFIG)
    
    @staticmethod
    def fetchall(sql, params=None):
        """执行查询并返回所有结果"""
        connection = None
        try:
            connection = DbOprate.get_connection()
            cursor = connection.cursor(pymysql.cursors.DictCursor)
            cursor.execute(sql, params or ())
            results = cursor.fetchall()
            return results
        except Exception as e:
            print(f"查询失败: {e}")
            return []
        finally:
            if connection:
                connection.close()
    
    @staticmethod
    def fetchone(sql, params=None):
        """执行查询并返回单个结果"""
        connection = None
        try:
            connection = DbOprate.get_connection()
            cursor = connection.cursor(pymysql.cursors.DictCursor)
            cursor.execute(sql, params or ())
            result = cursor.fetchone()
            return result
        except Exception as e:
            print(f"查询失败: {e}")
            return None
        finally:
            if connection:
                connection.close()
    
    @staticmethod
    def execute(sql, params=None):
        """执行SQL语句（INSERT/UPDATE/DELETE）"""
        connection = None
        try:
            connection = DbOprate.get_connection()
            cursor = connection.cursor()
            cursor.execute(sql, params or ())
            connection.commit()
            return cursor.rowcount
        except Exception as e:
            print(f"执行SQL失败: {e}")
            if connection:
                connection.rollback()
            return 0
        finally:
            if connection:
                connection.close()

    @staticmethod
    def save_result_to_mysql(result):
        """将单个结果保存到MySQL数据库"""
        connection = None
        try:
            connection = pymysql.connect(**DbOprate.DB_CONFIG)
            cursor = connection.cursor(pymysql.cursors.DictCursor)
            
            print(f"开始处理股票: {result['stock_code']}, 操作日期: {result['op_dt']}")
            
            # 获取最后一次记录
            select_sql = "SELECT * FROM stock_backtest WHERE strategy=%s AND stock_code=%s ORDER BY op_dt DESC LIMIT 1"
            print(f"执行SQL: {select_sql}")
            print(f"参数: {(result['strategy'], result['stock_code'])}")
            cursor.execute(select_sql, (result['strategy'], result['stock_code']))
            res_tmp = cursor.fetchone()
            print(f"历史记录查询结果: {res_tmp}")
            
            # 从传入的result中获取初始值，如果没有则从最后一次记录中获取
            if res_tmp:
                # 如果有历史记录，使用历史记录的值作为基础
                stock_num = float(res_tmp.get('stock_num', 0) or 0)
                stock_amount = float(res_tmp.get('stock_amount', 0) or 0)
                profitable = float(res_tmp.get('profitable', 0) or 0)
                can_use_money = float(res_tmp.get('can_use_money', 0) or 0) if res_tmp.get('can_use_money') is not None else 0
                init_cash = float(res_tmp.get('init_cash', 0) or 0) if res_tmp.get('init_cash') is not None else 0
            else:
                # 如果没有历史记录，使用传入的初始值
                stock_num = float(result.get('stock_num', 0) or 0)
                stock_amount = float(result.get('stock_amount', 0) or 0)
                profitable = float(result.get('profitable', 0) or 0)
                can_use_money = float(result.get('can_use_money', 100000) or 100000)  # 默认10万
                init_cash = float(result.get('init_cash', 100000) or 100000)  # 默认10万

            op_price = float(result['op_price'])
            op_num = int(result['op_num'])

            # 根据操作类型更新持仓和资金
            if result['op_act'] == 'sell':
                # 卖出：减少持仓，增加可用资金
                stock_num = float(stock_num - op_num)
                stock_amount = float(stock_amount - op_num * op_price)  # 持房价值减少
                can_use_money = float(can_use_money + op_num * op_price)  # 可用资金增加
                # 计算盈利
                if 'profit' in result:
                    profitable = float(profitable + result['profit'])
            
            elif result['op_act'] == 'buy':
                # 买入：增加持仓，减少可用资金
                stock_num = float(stock_num + op_num)
                stock_amount = float(stock_amount + op_num * op_price)  # 持仓价值增加
                can_use_money = float(can_use_money - op_num * op_price)  # 可用资金减少

            # 获取历史价格数据
            date_string = result['op_dt']
            dt = datetime.strptime(date_string, "%Y-%m-%d")
            op_dt = dt.timestamp()

            try:
                max_high_5d = 0.0
                min_low_5d = 0.0
                max_high_10d = 0.0
                min_low_10d = 0.0
                max_high_25d = 0.0
                min_low_25d = 0.0
            
                # 获取最近5天的数据，然后找最高价和最低价
                sql_5d = """
                SELECT high, low, trade_date FROM base_daily 
                WHERE ts_code=%s AND trade_date > %s 
                ORDER BY trade_date asc 
                LIMIT 5
                """
                cursor.execute(sql_5d, (result['stock_code'], op_dt))
                res_tmp_5d = cursor.fetchall()
                
                if res_tmp_5d and len(res_tmp_5d) > 0:
                    # 从最近5天的数据中找最高价和最低价
                    highs = [float(row['high']) for row in res_tmp_5d if row['high'] is not None and row['high'] != 0]
                    lows = [float(row['low']) for row in res_tmp_5d if row['low'] is not None and row['low'] != 0]
                    print(f"5日highs: {highs}")
                    print(f"5日lows: {lows}")
                    if highs and lows:
                        max_high_5d = max(highs)
                        min_low_5d = min(lows)
                    else:
                        max_high_5d = 0.0
                        min_low_5d = 0.0
                else:
                    max_high_5d = 0.0
                    min_low_5d = 0.0
                    print(f"⚠️ 5日数据为空")

                # 获取最近10天的数据，然后找最高价和最低价
                sql_10d = """
                SELECT high, low, trade_date FROM base_daily 
                WHERE ts_code=%s AND trade_date > %s 
                ORDER BY trade_date asc 
                LIMIT 10
                """
                cursor.execute(sql_10d, (result['stock_code'], op_dt))
                res_tmp_10d = cursor.fetchall()
                
                if res_tmp_10d and len(res_tmp_10d) > 0:
                    # 从最近10天的数据中找最高价和最低价
                    highs = [float(row['high']) for row in res_tmp_10d if row['high'] is not None and row['high'] != 0]
                    lows = [float(row['low']) for row in res_tmp_10d if row['low'] is not None and row['low'] != 0]
                    if highs and lows:
                        max_high_10d = max(highs)
                        min_low_10d = min(lows)
                    else:
                        max_high_10d = 0.0
                        min_low_10d = 0.0
                else:
                    max_high_10d = 0.0
                    min_low_10d = 0.0

                # 获取最近25天的数据，然后找最高价和最低价
                sql_25d = """
                SELECT high, low, trade_date FROM base_daily 
                WHERE ts_code=%s AND trade_date > %s 
                ORDER BY trade_date asc 
                LIMIT 25
                """
                cursor.execute(sql_25d, (result['stock_code'], op_dt))
                res_tmp_25d = cursor.fetchall()
                
                if res_tmp_25d and len(res_tmp_25d) > 0:
                    # 从最近10天的数据中找最高价和最低价
                    highs = [float(row['high']) for row in res_tmp_25d if row['high'] is not None and row['high'] != 0]
                    lows = [float(row['low']) for row in res_tmp_25d if row['low'] is not None and row['low'] != 0]
                    if highs and lows:
                        max_high_25d = max(highs)
                        min_low_25d = min(lows)
                    else:
                        max_high_25d = 0.0
                        min_low_25d = 0.0
                else:
                    max_high_25d = 0.0
                    min_low_25d = 0.0
                    
            except Exception as e:
                print(f"获取历史数据时出错: {e}")
                import traceback
                traceback.print_exc()
            # 判断当天是否增长小于5%
            sql_inc_flag = """
            SELECT open, close FROM base_daily WHERE ts_code=%s AND trade_date = %s 
            """
            cursor.execute(sql_inc_flag, (result['stock_code'], op_dt))
            flag_stock_tmp = cursor.fetchone()
            if flag_stock_tmp is not None and (flag_stock_tmp['close']/flag_stock_tmp['open'] < 1.05):
                # 插入数据
                insert_sql = """
                    INSERT INTO stock_backtest 
                    (strategy, stock_code, op_act, op_price, op_num, stock_num, stock_amount, profitable, op_dt, op_comm, max5, min5, max10, min10, init_cash, can_use_money,max25, min25)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """
                print(f"插入SQL: {insert_sql}")
                insert_params = (
                    result['strategy'],
                    result['stock_code'],
                    result['op_act'],
                    float(result['op_price']),
                    int(result['op_num']),
                    float(stock_num),
                    float(stock_amount),
                    float(profitable),
                    result['op_dt'],
                    float(result.get('op_comm', 0)),
                    float(max_high_5d) if max_high_5d is not None else 0.0,
                    float(min_low_5d) if min_low_5d is not None else 0.0,
                    float(max_high_10d) if max_high_10d is not None else 0.0,
                    float(min_low_10d) if min_low_10d is not None else 0.0,
                    float(init_cash),
                    float(can_use_money),
                    float(max_high_25d) if max_high_25d is not None else 0.0,
                    float(min_low_25d) if min_low_25d is not None else 0.0
                )
                print(f"插入参数: {insert_params}")

                cursor.execute(insert_sql, insert_params)
                
                connection.commit()
                print(f"✅ 结果已保存到数据库: {result['stock_code']}")
            
        except Exception as e:
            print(f"❌ 保存到数据库失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            if connection:
                connection.close()

    @staticmethod
    def save_result_to_mysql_all(result):
        """将单个结果保存到MySQL数据库"""
        print('ins db')
        connection = None
        try:
            connection = pymysql.connect(**DbOprate.DB_CONFIG)
            cursor = connection.cursor(pymysql.cursors.DictCursor)
            
            print(f"开始处理股票: {result['stock_code']}, 操作日期: {result['op_dt']}")
            
            # 获取最后一次记录
            select_sql = "SELECT * FROM stock_backtest WHERE strategy=%s AND stock_code=%s ORDER BY op_dt DESC LIMIT 1"
            print(f"执行SQL: {select_sql}")
            print(f"参数: {(result['strategy'], result['stock_code'])}")
            cursor.execute(select_sql, (result['strategy'], result['stock_code']))
            res_tmp = cursor.fetchone()
            print(f"历史记录查询结果: {res_tmp}")
            
            # 从传入的result中获取初始值，如果没有则从最后一次记录中获取
            if res_tmp:
                # 如果有历史记录，使用历史记录的值作为基础
                stock_num = float(res_tmp.get('stock_num', 0) or 0)
                stock_amount = float(res_tmp.get('stock_amount', 0) or 0)
                profitable = float(res_tmp.get('profitable', 0) or 0)
                can_use_money = float(res_tmp.get('can_use_money', 0) or 0) if res_tmp.get('can_use_money') is not None else 0
                init_cash = float(res_tmp.get('init_cash', 0) or 0) if res_tmp.get('init_cash') is not None else 0
            else:
                # 如果没有历史记录，使用传入的初始值
                stock_num = float(result.get('stock_num', 0) or 0)
                stock_amount = float(result.get('stock_amount', 0) or 0)
                profitable = float(result.get('profitable', 0) or 0)
                can_use_money = float(result.get('can_use_money', 100000) or 100000)  # 默认10万
                init_cash = float(result.get('init_cash', 100000) or 100000)  # 默认10万

            op_price = float(result['op_price'])
            op_num = int(result['op_num'])

            # 根据操作类型更新持仓和资金
            if result['op_act'] == 'sell':
                # 卖出：减少持仓，增加可用资金
                stock_num = float(stock_num - op_num)
                stock_amount = float(stock_amount - op_num * op_price)  # 持房价值减少
                can_use_money = float(can_use_money + op_num * op_price)  # 可用资金增加
                # 计算盈利
                if 'profit' in result:
                    profitable = float(profitable + result['profit'])
            
            elif result['op_act'] == 'buy':
                # 买入：增加持仓，减少可用资金
                stock_num = float(stock_num + op_num)
                stock_amount = float(stock_amount + op_num * op_price)  # 持仓价值增加
                can_use_money = float(can_use_money - op_num * op_price)  # 可用资金减少

            # 获取历史价格数据
            date_string = result['op_dt']
            dt = datetime.strptime(date_string, "%Y-%m-%d")
            op_dt = dt.timestamp()

            try:
                max_high_5d = 0.0
                min_low_5d = 0.0
                max_high_10d = 0.0
                min_low_10d = 0.0
                max_high_25d = 0.0
                min_low_25d = 0.0
            
                # 获取最近5天的数据，然后找最高价和最低价
                sql_5d = """
                SELECT high, low, trade_date FROM base_daily 
                WHERE ts_code=%s AND trade_date > %s 
                ORDER BY trade_date asc 
                LIMIT 5
                """
                cursor.execute(sql_5d, (result['stock_code'], op_dt))
                res_tmp_5d = cursor.fetchall()
                
                if res_tmp_5d and len(res_tmp_5d) > 0:
                    # 从最近5天的数据中找最高价和最低价
                    highs = [float(row['high']) for row in res_tmp_5d if row['high'] is not None and row['high'] != 0]
                    lows = [float(row['low']) for row in res_tmp_5d if row['low'] is not None and row['low'] != 0]
                    print(f"5日highs: {highs}")
                    print(f"5日lows: {lows}")
                    if highs and lows:
                        max_high_5d = max(highs)
                        min_low_5d = min(lows)
                    else:
                        max_high_5d = 0.0
                        min_low_5d = 0.0
                else:
                    max_high_5d = 0.0
                    min_low_5d = 0.0
                    print(f"⚠️ 5日数据为空")

                # 获取最近10天的数据，然后找最高价和最低价
                sql_10d = """
                SELECT high, low, trade_date FROM base_daily 
                WHERE ts_code=%s AND trade_date > %s 
                ORDER BY trade_date asc 
                LIMIT 10
                """
                cursor.execute(sql_10d, (result['stock_code'], op_dt))
                res_tmp_10d = cursor.fetchall()
                
                if res_tmp_10d and len(res_tmp_10d) > 0:
                    # 从最近10天的数据中找最高价和最低价
                    highs = [float(row['high']) for row in res_tmp_10d if row['high'] is not None and row['high'] != 0]
                    lows = [float(row['low']) for row in res_tmp_10d if row['low'] is not None and row['low'] != 0]
                    if highs and lows:
                        max_high_10d = max(highs)
                        min_low_10d = min(lows)
                    else:
                        max_high_10d = 0.0
                        min_low_10d = 0.0
                else:
                    max_high_10d = 0.0
                    min_low_10d = 0.0

                # 获取最近25天的数据，然后找最高价和最低价
                sql_25d = """
                SELECT high, low, trade_date FROM base_daily 
                WHERE ts_code=%s AND trade_date > %s 
                ORDER BY trade_date asc 
                LIMIT 25
                """
                cursor.execute(sql_25d, (result['stock_code'], op_dt))
                res_tmp_25d = cursor.fetchall()
                
                if res_tmp_25d and len(res_tmp_25d) > 0:
                    # 从最近10天的数据中找最高价和最低价
                    highs = [float(row['high']) for row in res_tmp_25d if row['high'] is not None and row['high'] != 0]
                    lows = [float(row['low']) for row in res_tmp_25d if row['low'] is not None and row['low'] != 0]
                    if highs and lows:
                        max_high_25d = max(highs)
                        min_low_25d = min(lows)
                    else:
                        max_high_25d = 0.0
                        min_low_25d = 0.0
                else:
                    max_high_25d = 0.0
                    min_low_25d = 0.0
                    
            except Exception as e:
                print(f"获取历史数据时出错: {e}")
                import traceback
                traceback.print_exc()
            # 插入数据
            insert_sql = """
                INSERT INTO stock_backtest 
                (strategy, stock_code, op_act, op_price, op_num, stock_num, stock_amount, profitable, op_dt, op_comm, max5, min5, max10, min10, init_cash, can_use_money,max25, min25)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            print(f"插入SQL: {insert_sql}")
            insert_params = (
                result['strategy'],
                result['stock_code'],
                result['op_act'],
                float(result['op_price']),
                int(result['op_num']),
                float(stock_num),
                float(stock_amount),
                float(profitable),
                result['op_dt'],
                float(result.get('op_comm', 0)),
                float(max_high_5d) if max_high_5d is not None else 0.0,
                float(min_low_5d) if min_low_5d is not None else 0.0,
                float(max_high_10d) if max_high_10d is not None else 0.0,
                float(min_low_10d) if min_low_10d is not None else 0.0,
                float(init_cash),
                float(can_use_money),
                float(max_high_25d) if max_high_25d is not None else 0.0,
                float(min_low_25d) if min_low_25d is not None else 0.0
            )
            print(f"插入参数: {insert_params}")

            cursor.execute(insert_sql, insert_params)
            
            connection.commit()
            print(f"✅ 结果已保存到数据库: {result['stock_code']}")
            
        except Exception as e:
            print(f"❌ 保存到数据库失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            if connection:
                connection.close()
